#include <windows.h>
#include <vector>
#include <fstream>
#include <sstream>
#include <string>
#include <list>

#include "../hwfs/hwfs_shared.h"
#include "debug.h"

#define NP_PROVIDER_CONTROL_KEY     L"SYSTEM\\CurrentControlSet\\Control\\NetworkProvider"

// splits string with NetworkProviders, delimited by ","
// skips HwFs provider
static void SplitProviders(const WCHAR* providerOrderStr, std::list<std::wstring>& providersList)
{
    providersList.clear();
    std::wstring providers(providerOrderStr);

    size_t pos = 0;
    while (1) {
        pos = providers.find(',');
        if (pos != std::wstring::npos) {
            std::wstring p = providers.substr(0, pos);
            if (p != HWFS_SERVICE_NAME)
                providersList.push_back(p);
            providers.erase(0, pos + 1);
        } else {
            providersList.push_back(providers);
            break;
        }
    }
}

static void ConcatProviders(const std::list<std::wstring>& providersList, std::wstring& result)
{
    result.clear();
    for (const std::wstring& s : providersList) {
        if (s.empty())
            continue;
        if (!result.empty())
            result += L",";
        result += s;
    }
}

static std::wstring GetHostsFileName()
{
    UINT MaxLen = MAX_PATH;
    while (1) {
        std::vector<WCHAR> WinDirBuf(MaxLen);
        UINT Len = GetSystemDirectoryW(WinDirBuf.data(), MaxLen);
        if (Len == 0) {
            DBG_PRINT((L"Failed to get System Dir: 0x%x\n", GetLastError()));
            return L"";
        }
        if (Len > MaxLen) {
            MaxLen = Len;
            continue;
        }
        std::wstring ret(WinDirBuf.data(), (size_t)Len);
        return ret + L"\\drivers\\etc\\hosts";
    }
}

typedef std::list<std::string> StringList;

static bool ReadEtcHosts(StringList& result)
{
    result.clear();
    std::wstring hosts = GetHostsFileName();
    if (hosts.empty()) {
        return false;
    }
    std::ifstream inp(hosts);
    if (!inp) {
        DBG_PRINT((L"Failed to open hosts file\n"));
        return false;
    }
    std::string str;
    while (std::getline(inp, str)) {
        result.push_back(str);
    }
    return true;
}

static bool WriteEtcHosts(StringList& content)
{
    std::wstring hosts = GetHostsFileName();
    if (hosts.empty()) {
        return false;
    }
    std::ofstream outp(hosts);
    if (!outp) {
        DBG_PRINT((L"Failed to open hosts file for write\n"));
        return false;
    }
    for (const std::string& st : content) {
        outp << st << std::endl;
    }
    return true;
}

static bool SplitHostsString(const std::string& s, std::string& t1, std::string& t2)
{
    auto iss = std::istringstream(s);
    std::string str;

    if (!(iss >> t1)) {
        return false;
    }

    if (t1[0] == '#') {
        return false;
    }

    if (!(iss >> t2)) {
        return false;
    }

    return true;
}

static void RemoveHwfsHostsEntry(StringList& content)
{
    for (auto it = content.begin(); it != content.end();) {
        std::string addr;
        std::string name;
        if (!SplitHostsString(*it, addr, name)) {
            ++it;
            continue;
        }
        if (0 == _stricmp(name.c_str(), "hwfs")) {
            content.erase(it++);
        } else {
            ++it;
        }
    }
}

static bool RegisterInHostsFile()
{
    StringList content;
    if (!ReadEtcHosts(content)) {
        return false;
    }

    RemoveHwfsHostsEntry(content);
    content.push_back("0.0.0.0        hwfs");

    return WriteEtcHosts(content);
}

static bool UnregisterInHostsFile()
{
    StringList content;
    if (!ReadEtcHosts(content)) {
        return false;
    }

    RemoveHwfsHostsEntry(content);
    return WriteEtcHosts(content);
}

extern "C"
DWORD APIENTRY
HWFSInstall()
{
    DBG_PRINT((L"HWFSInstall\n"));

    // Register itself in the registry
    HKEY hKey;
    LSTATUS Status = RegOpenKeyEx(
        HKEY_LOCAL_MACHINE, NP_PROVIDER_CONTROL_KEY, 0,
        KEY_QUERY_VALUE | KEY_WRITE, &hKey);
    if (0 != Status) {
        DBG_PRINT((L"Failed to open NetworkProvider key:%lu\n", Status));
        return Status;
    }

    WCHAR ProviderOrderStr[1024];
    DWORD Size = sizeof(ProviderOrderStr);
    Status = RegGetValue(
        hKey, L"Order", L"ProviderOrder",
        RRF_RT_REG_SZ, NULL, ProviderOrderStr, &Size);
    if (0 != Status) {
        DBG_PRINT((L"Failed to read ProviderOrder value:%lu\n", Status));
        ProviderOrderStr[0] = '\0';
    }

    std::list<std::wstring> providersList;
    SplitProviders(ProviderOrderStr, providersList);

    providersList.push_front(HWFS_SERVICE_NAME);

    std::wstring result;
    ConcatProviders(providersList, result);

    Status = RegSetKeyValue(
        hKey, L"Order", L"ProviderOrder",
        REG_SZ, result.c_str(), (DWORD)(sizeof(WCHAR) * (1 + result.size())));
    if (0 != Status) {
        DBG_PRINT((L"Failed to register in ProviderOrder key:%lu\n", Status));
        RegCloseKey(hKey);
        return Status;
    }

    RegCloseKey(hKey);

    if (!RegisterInHostsFile()) {
        DBG_PRINT((L"Failed to register in etc/hosts\n"));
    }

    return 0;
}

extern "C"
DWORD APIENTRY
HWFSUninstall()
{
    DBG_PRINT((L"HWFSUninstall\n"));

    if (!UnregisterInHostsFile()) {
        DBG_PRINT((L"Failed to unregister in etc/hosts\n"));
    }

    // Register itself in the registry
    HKEY hKey;
    LSTATUS Status = RegOpenKeyEx(
        HKEY_LOCAL_MACHINE, NP_PROVIDER_CONTROL_KEY, 0,
        KEY_QUERY_VALUE | KEY_WRITE, &hKey);
    if (0 != Status) {
        DBG_PRINT((L"Failed to open NetworkProvider key:%lu\n", Status));
        return Status;
    }

    WCHAR ProviderOrderStr[1024];
    DWORD Size = sizeof(ProviderOrderStr);
    Status = RegGetValue(
        hKey, L"Order", L"ProviderOrder",
        RRF_RT_REG_SZ, NULL, ProviderOrderStr, &Size);
    if (0 != Status) {
        DBG_PRINT((L"Failed to read ProviderOrder value:%lu\n", Status));
        ProviderOrderStr[0] = '\0';
    }

    std::list<std::wstring> providersList;
    SplitProviders(ProviderOrderStr, providersList);

    std::wstring result;
    ConcatProviders(providersList, result);

    Status = RegSetKeyValue(
        hKey, L"Order", L"ProviderOrder",
        REG_SZ, result.c_str(), (DWORD)(sizeof(WCHAR) * (1 + result.size())));
    if (0 != Status) {
        DBG_PRINT((L"Failed to unregister in ProviderOrder key:%lu\n", Status));
        RegCloseKey(hKey);
        return Status;
    }

    RegCloseKey(hKey);

    return 0;
}
